#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <vector>
#include <chrono>

namespace py = pybind11;

// 矩陣乘法 (Naive + 多執行緒可以加強)
py::array_t<float> matmul(const py::array_t<float>& a, const py::array_t<float>& b) {
    auto buf_a = a.request(), buf_b = b.request();

    if (buf_a.ndim != 2 || buf_b.ndim != 2 || buf_a.shape[1] != buf_b.shape[0])
        throw std::runtime_error("形狀不符合乘法規則");

    ssize_t M = buf_a.shape[0];
    ssize_t K = buf_a.shape[1];
    ssize_t N = buf_b.shape[1];

    auto result = py::array_t<float>({M, N});
    auto buf_r = result.mutable_unchecked<2>();

    auto ptr_a = static_cast<float*>(buf_a.ptr);
    auto ptr_b = static_cast<float*>(buf_b.ptr);

    for (ssize_t i = 0; i < M; ++i) {
        for (ssize_t j = 0; j < N; ++j) {
            float sum = 0.0f;
            for (ssize_t k = 0; k < K; ++k) {
                sum += ptr_a[i * K + k] * ptr_b[k * N + j];
            }
            buf_r(i, j) = sum;
        }
    }
    return result;
}